백준 13537번

수열과 쿼리 1

Posted by ChoiCube84 on August 28, 2023 · 16 mins read

백준 13537번

오늘 풀어본 문제는 백준의 13537번 문제1이다. 문제 풀이에 사용한 언어는 C++ 이다.

solved.ac 기준 CLASS

CLASS 7

문제 정보

이 문제의 내용과 조건은 다음과 같다.

문제

길이가 $N$ 인 수열 $A_1$, $A_2$, …, $A_N$ 이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

  • $i\ j\ k$ : $A_i$, $A_{i+1}$, …, $A_j$ 로 이루어진 부분 수열 중에서 $k$ 보다 큰 원소의 개수를 출력한다.

입력

첫째 줄에 수열의 크기 $N$ $(1 \leq N \leq 100,000)$ 이 주어진다.

둘째 줄에는 $A_1$, $A_2$, …, $A_N$ 이 주어진다. $(1 \leq A_i \leq 109)$

셋째 줄에는 쿼리의 개수 $M$ $(1 \leq M \leq 100,000)$ 이 주어진다.

넷째 줄부터 M개의 줄에는 쿼리 $i,\ j,\ k$ 가 한 줄에 하나씩 주어진다. $(1 \leq i \leq j \leq N,\ 1 \leq k \leq 109)$

출력

각각의 쿼리마다 정답을 한 줄에 하나씩 출력한다.

풀이과정

1번째 시도

이 문제는 머지소트 트리를 활용하는 문제이다. 세그먼트 트리와 비슷한 부분이 있는데, 세그먼트 트리의 각 노드가 특정 구간에서의 특정 결과를 나타내듯이, 머지소트 트리는 각 노드가 특정 구간의 요소들이 정렬된 벡터를 가진다고 생각하면 된다.

내가 떠올린 해결방법은 머지소트 트리를 구성하고 나서 특정 구간의 쿼리로 그 구간의 정렬된 배열을 받아오면, std::upper_bound 함수를 활용하여 쿼리로 주어진 값보다 큰 값의 인덱스를 알아내어 전체 개수에서 빼는 방식으로 문제를 해결하는 것이었다.

코드는 다음과 같이 작성하였다.

#include <bits/stdc++.h>

using namespace std;

#ifndef __MERGE_SORT_TREE__
#define __MERGE_SORT_TREE__

#include <vector>
#include <algorithm>

template <typename T>
class MergeSortTree {
private:
    std::vector<std::vector<T>> tree;
    std::vector<T> arr;
    size_t n;

    void build(size_t idx, size_t left, size_t right) {
        if (left == right) {
            tree[idx].resize(1);
            tree[idx][0] = arr[left];
            return;
        }

        size_t mid = (left + right) / 2;

        build(2 * idx, left, mid);
        build(2 * idx + 1, mid + 1, right);

        tree[idx].resize(tree[2 * idx].size() + tree[2 * idx + 1].size());
        std::merge(tree[2 * idx].begin(), tree[2 * idx].end(),
            tree[2 * idx + 1].begin(), tree[2 * idx + 1].end(),
            tree[idx].begin());
    }

    std::vector<T> query(size_t i, size_t j, size_t node, size_t currentLeft, size_t currentRight) {
        std::vector<T> result;

        if (j < currentLeft || currentRight < i) {
            // Do nothing
        }
        else if (i <= currentLeft && currentRight <= j) {
            result = tree[node];
        }
        else {
            size_t mid = (currentLeft + currentRight) / 2;
            const std::vector<T>& leftArray = query(i, j, 2 * node, currentLeft, mid);
            const std::vector<T>& rightArray = query(i, j, 2 * node + 1, mid + 1, currentRight);

            result.resize(leftArray.size() + rightArray.size());
            std::merge(leftArray.begin(), leftArray.end(), rightArray.begin(), rightArray.end(), result.begin());
        }

        return result;
    }

public:
    MergeSortTree(const std::vector<T>& a) : arr(a), n(a.size()) {
        tree.resize(4 * n + 1);
        build(1, 0, n - 1);
    }

    std::vector<T> query(size_t i, size_t j) {
        return query(i, j, 1, 0, n - 1);
    }
};

#endif // !__MERGE_SORT_TREE__


int main(void) {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	int N, M;
	cin >> N;

	vector<int> a(N);
	for (int i = 0; i < N; i++) {
		cin >> a[i];
	}

	MergeSortTree<int> mst(a);

	cin >> M;

	for (int query = 0; query < M; query++) {
		int i, j, k;
		cin >> i >> j >> k;

		vector<int> sortedArray = mst.query(i - 1, j - 1);
		cout << sortedArray.size() - (upper_bound(sortedArray.begin(), sortedArray.end(), k) - sortedArray.begin()) << "\n";
	}

	return 0;
}

실행 결과 ‘시간 초과’가 나왔다.

2번째 시도

코드의 시간 복잡도에는 문제가 없어보였기 때문에 벡터를 반환하는 과정에서 불필요한 복사가 이루어지지 않도록 수정하였다. 또한 쿼리 함수를 좀 더 사용하기 편하도록 main 함수에서 개수를 알아내는 대신 쿼리 함수 내부적으로 개수를 알아낸 뒤 그 값을 반환하도록 수정하였다.

코드는 다음과 같이 작성하였다.

#include <bits/stdc++.h>

using namespace std;

template <typename T>
class MergeSortTree {
private:
    std::vector<std::vector<T>> tree;
    std::vector<T> arr;
    size_t n;

    void build(size_t idx, size_t left, size_t right) {
        if (left == right) {
            tree[idx].resize(1);
            tree[idx][0] = arr[left];
            return;
        }

        size_t mid = (left + right) / 2;

        build(2 * idx, left, mid);
        build(2 * idx + 1, mid + 1, right);

        tree[idx].resize(tree[2 * idx].size() + tree[2 * idx + 1].size());
        std::merge(tree[2 * idx].begin(), tree[2 * idx].end(),
            tree[2 * idx + 1].begin(), tree[2 * idx + 1].end(),
            tree[idx].begin());
    }

    size_t query(size_t i, size_t j, const T& k, size_t node, size_t currentLeft, size_t currentRight) {
        size_t result;

        if (j < currentLeft || currentRight < i) {
            result = 0;
        }
        else if (i <= currentLeft && currentRight <= j) {
            result = tree[node].end() - upper_bound(tree[node].begin(), tree[node].end(), k);
        }
        else {
            size_t mid = (currentLeft + currentRight) / 2;
            result = query(i, j, k, 2 * node, currentLeft, mid) + query(i, j, k, 2 * node + 1, mid + 1, currentRight);
        }

        return result;
    }

public:
    MergeSortTree(const std::vector<T>& a) : arr(a), n(a.size()) {
        tree.resize(4 * n + 1);
        build(1, 0, n - 1);
    }

    size_t query(size_t i, size_t j, const T& k) {
        return query(i, j, k, 1, 0, n - 1);
    }
};

int main(void) {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	int N, M;
	cin >> N;

	vector<int> a(N);
	for (int i = 0; i < N; i++) {
		cin >> a[i];
	}

	MergeSortTree<int> mst(a);

	cin >> M;

	for (int query = 0; query < M; query++) {
		int i, j, k;
		cin >> i >> j >> k;

        cout << mst.query(i - 1, j - 1, k) << "\n";
	}

	return 0;
}

그러자 모든 테스트 케이스를 통과하고 정답이 나오는 것을 확인할 수 있었다.

마무리

이번 문제에서 사용된 자료구조인 머지소트 트리도 세미나에서 배운 내용이었다. 요즘 CLASS 4나 5에 있는 문제보다도 6, 7, 8에 있는 문제들에 더 관심이 가는 것 같다.

아무래도 요즘 내가 레이팅에 신경을 쓰고 있는 것 같다. 자료구조를 한 번 만들어 두기만 하면 몇 번이고 사용가능하기 때문에 이를 이용하여 문제들을 ‘날먹’하여 빠르게 레이팅을 올린 심산인 것이다. 하지만 이건 내가 처음에 PS를 시작하는 초심을 잃어가는 현상일지도 모르겠다.

처음부터 레이팅을 올리는게 목표였다면 클래스 문제들을 차근차근 풀어볼 것이 아니라 잘하는 주제를 하나 잡아서 그런 문제들만 주구장창 푸는 것이 더 빨랐을 것이다. 하지만 굳이 클래스 문제들을 푼 것은 PS에 필요한 ‘여러 가지’ 기법들을 배우고, ‘사고력’을 기르기 위해서였다. 그렇게 생각해보면 레이팅을 신경쓰는 건 클래스 문제들을 푸는 근본적인 이유와는 상충되는 것이다.

물론 레이팅에 마음이 갈 수도 있고 신경을 안 쓰기는 어렵지만, 그것에 매몰되어 초심을 잃어서는 안된다고 생각한다. 앞으로 내가 꿋꿋이 PS를 해나가려면 마음을 다시 굳게 다잡아야겠다.

참고로 오늘 구현한 머지소트 트리 자료구조는 레포지토리2에 저장해두었다. 필요한 사람은 링크를 타고 들어가 확인해보면 된다.

오늘의 PS는 여기까지!


1: https://www.acmicpc.net/problem/13537
2: https://github.com/ChoiCube84/baekjoon-solutions/blob/main/C%2B%2B/custom_data_structures/merge_sort_tree.hpp